import argparse
import json
import os
import re
from subprocess import Popen
import subprocess
import time
import codecs
from openpyxl import Workbook
from openpyxl.styles import PatternFill
from tqdm.contrib.concurrent import thread_map


def run_dir(currdir, test_list, basename='', maxlevel=1000):
    dir_list = os.listdir(currdir)
    CONFIG_JSON_NAME = 'config.json'
    if CONFIG_JSON_NAME in dir_list:
        # with json configuration file
        with open(os.path.join(currdir, CONFIG_JSON_NAME), 'r') as conf:
            conf = json.load(conf)

        name = os.path.basename(currdir)
        if 'name' in conf.keys():
            name = conf['name']

        if 'maxlevel' in conf.keys():
            maxlevel = min(maxlevel, int(conf['maxlevel']))

        if 'import' in conf.keys():
            # import other dir
            for i in conf['import']:
                d = os.path.join(currdir, i)
                if not os.path.isdir(d):
                    continue  # is not a dir
                if basename:
                    n = '%s.%s' % (basename, name)
                else:
                    n = name
                # print("<import>", d, ',', n, ',', maxlevel)
                run_dir(d, test_list, n, maxlevel)
            pass

        if 'case' in conf.keys():
            if basename:
                n = '%s.%s' % (basename, name)
            else:
                n = name
            for item in conf['case']:
                if int(item[1]) <= maxlevel:
                    add_case(n, os.path.join(currdir, item[0]), test_list)

        pass
    elif maxlevel >= 0:
        # only testcases in dir
        c = get_all_case_filename(currdir)
        for item in c:
            add_case(
                '%s.%s' % (basename, os.path.basename(currdir)),
                os.path.join(currdir, item),
                test_list
            )
    return


def add_case(testname, casedir, test_list):
    # print('<case>', testname, '@', casedir)
    if testname not in test_list.keys():
        test_list[testname] = list()
    test_list[testname].append(casedir)
    return


def get_all_case_filename(dir):
    return list(
        set(
            os.path.splitext(casename)[0]
            for casename in os.listdir(dir)
            if os.path.splitext(casename)[1] in ['.in', '.out']
        )
    )


def check(case):
    '''
        case[0] = test name
        case[1] = testcase path
    '''
    with codecs.open(os.path.join(basedir, case[1]+'.in'), 'r', encoding='utf-8') as f:
        testcase_input = f.read()
    with codecs.open(os.path.join(basedir, case[1]+'.out'), 'r', encoding='utf-8') as f:
        testcase_answer = f.read()

    # print(testcase_input)
    testcase_input = testcase_input.lower()
    inputs = testcase_input.splitlines()
    testcase_input = []
    # print(inputs)
    for input in inputs:
        # print(input)
        input = re.sub(r'(#|\/\/).*$', ' ', input)
        input = re.sub(r'\s+', ' ', input.strip())
        if input.startswith('preset seq '):
            input = input.replace('preset seq ', '')
            input = re.sub(r'q', '1', input)
            input = re.sub(r'a', '2', input)
            input = re.sub(r's', '3', input)
            input = re.sub(r'j', '4', input)
            input = input + '\n10000'
        testcase_input.append(input)
        # print(input)
    testcase_input = '\n'.join(testcase_input)

    # print(testcase_input)
    output = ''
    EXPIRE_TIME_MS = 5000
    proc = Popen(
        args.name, shell=True, universal_newlines=True,
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        encoding='utf-8'
    )
    try:
        t = time.time_ns()
        _, output = proc.communicate(
            testcase_input + '\nprint\nquit\n',
            timeout=EXPIRE_TIME_MS/1000)
        t = time.time_ns() - t
    except subprocess.TimeoutExpired:
        t = time.time_ns() - t
        print('Timeout:', case[1])
        os.system('taskkill /F /T /PID %i' % proc.pid)

    # print(output)
    out_kv = parse_key_value_pairs(output.splitlines()) if output else set()
    ans_kv = parse_key_value_pairs(testcase_answer.splitlines())
    # print(out_kv, ans_kv)
    # if out_kv != ans_kv:
    #     print(out_kv, ans_kv)

    return (
        case[0],
        os.path.basename(case[1]),
        out_kv == ans_kv,
        t / 1e6,
        '\n'.join(ans_kv ^ out_kv) if out_kv != ans_kv else '',
    )


def parse_key_value_pairs(lines):
    ret = []
    kv_pat = re.compile(
        r'([a-zA-Z]+)(?:\s+(\w+))(?:\s+(\w+))?(?:\s+(\w+))?(?:\s+(\w+))?(?:\s+(\w+))?(?:\s+(\w+))?(?:\s+(\w+))?(?:\s+(\w+))?')
    for line in lines:
        line = trim_line(line)
        if not line:
            continue
        # print(line)
        mat = kv_pat.match(line)
        if mat:
            ret.append(' '.join((e for e in mat.groups() if e)))
        else:
            print('Warn::', line)
            # assert False
    # print(ret)
    ret.sort()
    ret = set(ret)
    return ret


def trim_line(line):
    if not line:
        return ''

    line = re.sub(r'(#|\/\/).*$', ' ', line)  # comment
    line = re.sub(r'\s+', ' ', line.strip())  # strip and merge empty char
    if not line:
        return ''
    line = line.lower()  # lowercase
    CAPTURE_CMD = ('seq', 'q', 'a', 's', 'j', 'building')
    if not any(line.startswith(c + ' ') for c in CAPTURE_CMD):
        return ''  # check first word

    return line


def export_to_xlsx(file_name, res):
    wb = Workbook()
    ws = wb.active  # work sheet
    cnt_all = 0
    all = 0
    row = 0
    res_keys = list(res.keys())
    res_keys.sort()
    for test in res_keys:
        cnt = 0
        for item in res[test]:
            if item[2]:
                cnt = cnt + 1  # count True
        cnt_all = cnt_all + cnt
        all = all + len(res[test])
        row = row + 1
        ws.cell(row, 1).value = 'testname'
        ws.cell(row, 2).value = test
        ws.cell(row, 3).value = '%d/%d' % (cnt, len(res[test]))

        row = row + 1
        ws.cell(row, 1).value = 'testname'
        ws.cell(row, 2).value = 'testcase'
        ws.cell(row, 3).value = 'status'
        ws.cell(row, 4).value = 'time(ms)'
        ws.cell(row, 5).value = 'diff'

        GREEN = PatternFill(start_color='7FFF00',
                            end_color='7FFF00', fill_type='solid')
        RED = PatternFill(start_color='FF0000',
                          end_color='FF0000', fill_type='solid')
        EXPIRE_TIME_MS = 5000
        for item in res[test]:
            row = row + 1
            ws.cell(row, 1).value = item[0]
            ws.cell(row, 2).value = item[1]
            ws.cell(row, 3).value = 'PASS' if item[2] else 'FAIL'
            ws.cell(row, 3).fill = GREEN if item[2] else RED
            ws.cell(row, 4).value = item[3]
            if item[3] >= EXPIRE_TIME_MS:
                ws.cell(row, 4).fill = RED
            ws.cell(row, 5).value = item[4]
            pass  # end for

        row = row + 1  # empty line
        pass
    row = row + 1
    ws.cell(row, 1).value = 'TOTAL'
    ws.cell(row, 2).value = 'RESULT'
    ws.cell(row, 3).value = '%d/%d' % (cnt_all, all)
    wb.save(file_name)
    return


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--dir', help='testcase dir',
                        default='.', type=str)
    parser.add_argument('-w', '--workers',
                        help='numebr of work processes', default=20, type=int)
    parser.add_argument('-n', '--name', help='process name',
                        default='..\\..\\game\\build\\rich.exe', type=str)
    parser.add_argument(
        '-e', '--export', help='export to file', default=True, type=bool)
    parser.add_argument('-f', '--exportfile',
                        help='export file name', default='report.xlsx', type=str)
    args = parser.parse_args()

    basedir = args.dir
    test_list = dict()

    run_dir(basedir, test_list)
    # print(test_list)

    testcases = list()
    keys = list(test_list.keys())
    keys.sort()
    for k in keys:
        for v in test_list[k]:
            testcases.append((k, v))
    testcases = set(testcases)

    print('run', args.name)
    res = thread_map(check, testcases, max_workers=args.workers)

    test_res = dict()
    for item in res:
        if item[0] not in test_res:
            test_res[item[0]] = list()
        test_res[item[0]].append(item)
    for item in test_res:
        test_res[item].sort(key=lambda x: x[1])
        # print(test_res)

    if args.export:
        export_to_xlsx(args.exportfile, test_res)
        print('export to', args.exportfile)
